import torch
from scipy.io import loadmat


def get_device_str() -> str:
    if hasattr(torch, "cuda") and torch.cuda.is_available():
        return "cuda"
    elif hasattr(torch, "xpu") and torch.xpu.is_available():
        return "xpu"
    elif hasattr(torch, "mps") and torch.mps.is_available():
        return "mps"
    else:
        return "cpu"


class HARDatasetPt(torch.utils.data.Dataset):
    def __init__(self, data_path: str, x_key: str = "train_X", y_key: str = "train_y",*args, **kwargs):
        super(HARDatasetPt, self).__init__(*args, **kwargs)
        self.x_key = x_key
        self.y_key = y_key
        self.data = torch.load(data_path)
        self.data[self.y_key] = self.data[self.y_key].squeeze()
        assert len(self.data[self.x_key]) == len(self.data[self.y_key]), f"{len(self.data[self.x_key])} != {len(self.data[self.y_key])}"

    def __len__(self):
        return len(self.data[self.x_key])

    def __getitem__(self, index):
        return self.data[self.x_key][index], self.data[self.y_key][index]


class HARDatasetMat(torch.utils.data.Dataset):
    def __init__(self, data_path: str, x_key: str = "test_X", y_key: str = "test_y",*args, **kwargs):
        super(HARDatasetMat, self).__init__(*args, **kwargs)
        self.x_key = x_key
        self.y_key = y_key
        self.data = loadmat(data_path)
        self.data[self.x_key] = torch.from_numpy(self.data[self.x_key])
        self.data[self.y_key] = torch.from_numpy(self.data[self.y_key].squeeze() - 1)
        assert len(self.data[self.x_key]) == len(self.data[self.y_key]), f"{len(self.data[self.x_key])} != {len(self.data[self.y_key])}"

    def __len__(self):
        return len(self.data[self.x_key])

    def __getitem__(self, index):
        return self.data[self.x_key][index], self.data[self.y_key][index]
